#!/usr/bin/python
# Author: Jimmy Saw
#
#                  91       100
#                  CTG      ATG         TGA
# ---[  X(-) ]--------------[    Y(+)     ]--------[      +      ]
# IGS|       |       IGS    |             |   IGS  |
#    TTA   CTA         CAA     
#    {3}{2}{1}              {1}{2}{3}{4}{5} 
#           85          97
#
# When fixing Y(+) which is in plus strand,
# this script will look for start codons upstream IGS region (left of ATG at 100)
# If found, it will add new bases that extend to left of original start site.
# New start will then be 91 instead of 100.
# When fixing X(-) which is in minus strand, upstream region is actually on right.
# If found, it will add new bases that extend to right of original start site.
# New start will then be 97 instead of 85.
# -------------IMPORTANT!!!!!--------------------------------------------------
# You are actually giving it coordinates of coding regions
# Example: If working on Y(+), gene start is 100 and gene stop is 115.

import sys
import string
from Bio import SeqIO

sf = open(sys.argv[1], "rU")

record = SeqIO.read(sf, "fasta")

start = raw_input("Enter gene start: ")

stop = raw_input("Enter gene stop: ")

strand = raw_input("Enter strand (+ or -): ")

fixtype = raw_input("cut or extend? ")

fixlength_aa = raw_input("How many residues (amino acid) to extend or cut? ")

fixlength = int(fixlength_aa) * 3

start_codons = {'ATG': 'Start', 'GTG': 'Start', 'TTG': 'Start'}
stop_codons = {'TAA': 'Stop', 'TAG': 'Stop', 'TGA': 'Stop'}

def getstart(triplist):
    extension = 0 
    if len(triplist) > 0:
        for index, item in enumerate(triplist):
            codon = str(item)
            if codon not in stop_codons:
                if codon in start_codons:
                    extension = (index + 1) * 3
                    break
                else:
                    extension = 0
                    index = 0
                    pass
            else:
                extension = 0
                index = 0
                break
    else:
        extension = 0
        index = 0

    return extension, index 

new_start = 0

n = 3

triplets = []

ext = 0
cut = 0
ind = 0

if strand == "+":
    if fixtype == "extend":
        substart = int(start) - fixlength
        substop = int(start) - 1
        subseq = record.seq[substart:substop]
        num = len(subseq)
        print "Substring = ", subseq, num, "bp"
        while n <= num:
            t = subseq[num-n:num-n+3]
            triplets.append(t)
            n += 3
        ext, ind = getstart(triplets)
        new_start = int(start) - ext
    else:
        substart = int(start) + 3 #skip the first start codon that's already identified
        substop = substart + fixlength
        subseq = record.seq[substart:substop]
        num = len(subseq)
        print "Substring = ", subseq, num, "bp"
        while n <= num:
            t = subseq[0+n-3:0+n]
            triplets.append(t)
            n += 3
        cut, ind = getstart(triplets)
        new_start = int(start) + cut

else:
    if fixtype == "extend":
        substart = int(stop) + 1
        substop = int(stop) + fixlength
        subseq = record.seq[substart:substop].reverse_complement()
        num = len(subseq)
        print "Substring = ", subseq, num, "bp"
        while n <= num:
            t = subseq[num-n:num-n+3]
            triplets.append(t)
            n += 3
        ext, ind = getstart(triplets)
        new_start = int(stop) + ext
    else:
        substart = int(stop) - 3 - fixlength #skip the first start codon that's already identified
        substop = int(stop) - 3 #1st 3 bases of start codon in minus strand skipped
        subseq = record.seq[int(start) - fixlength:int(start)].reverse_complement()
        num = len(subseq)
        print "Substring = ", subseq, num, "bp"
        while n <= num:
            t = subseq[0+n-3:0+n]
            triplets.append(t)
            n += 3
        cut, ind = getstart(triplets)
        new_start = int(stop) - cut

#print "\n"
#print "Start codon found: "

outstring = ""
codon_found = ""
index_found = 0

if len(triplets) > 0:
    for index, item in enumerate(triplets):
        codon = str(item)
        if codon not in stop_codons:
            if codon in start_codons:
                codon_found = codon
                index_found = index
                break
            else:
                codon_found = "Not found"
                index_found = 0
                pass
        else:
            codon_found = "Not found"
            index_found = 0
            break
else:
    outstring = "Region to fix is less than 3bp long!"

#outstring = str(index_found) + "\t" + codon_found 

#print outstring

print "Start codon found: " + codon_found

print "New Start location:" + "\t" + str(new_start)

print "Triplets: ",

if fixtype == "extend":
    m = -1 
    neg = -len(triplets)
    while m >= neg:
        fix = str(triplets[ind])
        triplets[ind] = fix.lower()
        print triplets[m],
        m -= 1
    print "\n"
else:
    m = 0
    pos = len(triplets)
    while m < pos:
        fix = str(triplets[ind])
        triplets[ind] = fix.lower()
        print triplets[m],
        m += 1
    print "\n"

sf.close()
